[serge] Fix 2 integration tests for model cohere2_vision failing with OOM, raising OutOfMemoryError (other (2)) - #50483

Original CI failure

• Failure group: 2 integration tests for model cohere2_visionfailing withOOM, raising OutOfMemoryError (other (2))
• tests/models/cohere2_vision/test_modeling_cohere2_vision.py::Cohere2IntegrationTest::test_model_integration_batched_generate — retained memory (fixable): needed 288 MiB, PyTorch already held 21.68 GiB of 22.30 GiB

(line 1292)  torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 288.00 MiB. GPU 1 has a total capacity of 22.30 GiB of which 106.69 MiB is free. Process 121671 has 22.19 GiB memory in use. Of the allocated memory 21.68 GiB is allocated by PyTorch, and 273.84 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)

Where to watch it:

• test_model_integration_batched_generate — Test dashboard

Relates to #50482

test_model_integration_forward and test_model_integration_batched_generate OOM on a 288 MiB copy during from_pretrained.

Not retained memory: the class already uses MemoryCleanupMixin, each test loads exactly one model, and both reproduced tracebacks fail at the same weight key with the same 21.68 GiB already allocated on GPU 1 — that is the checkpoint being loaded in that call, not a leftover from an earlier test.

Root cause is the load pattern: device_map="auto" budgets each device to its full capacity when several are visible, so the last weight copy has no headroom. The patch adds the per-GPU max_memory cap (70%, plus a CPU budget) already used by phimoe, glm4_moe, glm4v_moe, flex_olmo and qwen3_omni_moe for the same failure. No coverage is reduced and no expected values change.

✅ Verified on GPU

serge ran the targeted @slow test(s) on a GPU runner and opened this PR only after they passed with this patch. Each targeted test was run 5× on both the pre-patch and patched trees to rule out flakiness — the result held on every run.

• Passes with this patch: https://github.com/huggingface/transformers/actions/runs/42532460106
• Failed before the fix (base commit): https://github.com/huggingface/transformers/actions/runs/42529843470

Possibly related

Existing issues/PRs mentioning test_model_integration_batched_generate (keyword match — not verified to share a root cause):

• #50430 — Fix! (PR, closed, updated 2026-11-06)
• #48589 — Synthetic test assets (PR, closed, updated 2026-09-10)
• #45280 — add Qianfan-OCR model definition (PR, closed, updated 2026-04-17)

This change was produced automatically by serge from a CI failure report. The patch was generated by an LLM and applied by serge; review before merging.

serge v0.1.0 · model: deepseek-ai/DeepSeek-V4.1-Flash-1017:together · 46 LLM turns · 69 tool calls · 585.3s · 2157099 in / 203496 out tokens